import math
import time
from typing import Dict, Any

import numpy as np
import pandas as pd
import sympy
from scipy.optimize import root


class JMMethod:

    def __init__(self, failure_dataframe, precision):
        self.dL_N0_dN0 = None
        self.dL_N0_dphi = None
        self.dL_phi_dphi = None
        self.dL_phi_dN0 = None
        self.L_N0 = None
        self.L_phi = None
        self.__para = {"phi": None, "N0": None}
        self.__fail = failure_dataframe["value"].tolist()
        self.__precision = precision

    @staticmethod
    def get_para(phi, N0) -> Dict[str, Any]:
        return {"phi": phi, "N0": N0}

    @staticmethod
    def get_numerical_Z(para, i) -> float:
        phi = para["phi"]
        N0 = para["N0"]
        return phi * (N0 - i + 1)

    @staticmethod
    def get_numerical_f(para, x_i, i) -> float:
        return JMMethod.get_numerical_Z(para, i) * math.exp(-JMMethod.get_numerical_Z(para, i) * x_i)

    @staticmethod
    def get_numerical_F(para, x_i, i) -> float:
        return 1 - math.exp(-JMMethod.get_numerical_Z(para, i) * x_i)

    @staticmethod
    def get_numerical_R(para, x_i, i) -> float:
        return 1 - JMMethod.get_numerical_F(para, x_i, i)

    @staticmethod
    def get_symbol_Z() -> Any:
        phi, N0, i = sympy.symbols("phi N0 i", positive=True)
        return phi * (N0 - i + 1)

    @staticmethod
    def get_symbol_f() -> Any:
        x = sympy.symbols("x", positive=True)
        return JMMethod.get_symbol_Z() * JMMethod.get_symbol_R()

    @staticmethod
    def get_symbol_R() -> Any:
        x = sympy.symbols("x", positive=True)
        return sympy.exp(-JMMethod.get_symbol_Z() * x)

    @staticmethod
    def get_symbol_F() -> Any:
        x = sympy.symbols("x", positive=True)
        return 1 - JMMethod.get_symbol_R()

    @staticmethod
    def loading(drop_frame) -> Any:
        time.sleep(drop_frame)

    def get_MLE(self) -> Any:
        phi, N0, i, x, n = sympy.symbols("phi N0 i x n", positive=True)
        fail_list = self.__fail[1: 5]
        L = np.prod([self.get_symbol_f().subs({x: item, i: fail_list.index(item)}) for item in fail_list])
        self.L_phi = sympy.diff(sympy.expand_log(sympy.log(L)), phi)
        self.L_N0 = sympy.diff(sympy.expand_log(sympy.log(L)), N0)
        self.dL_phi_dphi = sympy.diff(self.L_phi, phi)
        self.dL_phi_dN0 = sympy.diff(self.L_phi, N0)
        self.dL_N0_dphi = sympy.diff(self.L_N0, phi)
        self.dL_N0_dN0 = sympy.diff(self.L_N0, N0)

    def func_index(self, x):
        phi, N0 = sympy.symbols("phi, N0", position=True)
        para = {phi: x[0], N0: x[1]}
        func = [self.L_phi.subs(para), self.L_N0.subs(para)]
        differ_func = [[self.dL_phi_dphi.subs(para), self.dL_phi_dN0.subs(para)],
                       [self.dL_N0_dphi.subs(para), self.dL_N0_dN0.subs(para)]]
        return func, differ_func

    def get_para_estimate(self):
        sol = root(self.func_index, [1, 1], jac=True, method="lm")
        print(sol.x)


if __name__ == "__main__":
    df = pd.DataFrame({"value": [1, 2, 3, 5, 5]})
    JMmethod = JMMethod(df, None)
    JMmethod.get_MLE()
